package proxy

import (
	"bytes"
	"com.yy.api/my_gateway/internal/types/models"
	"io"
	"net/http"
	"net/http/httputil"
	"net/url"
)

func newChatProxy(apiKey, baseUrl string) *httpProxyImpl {
	target, _ := url.Parse(baseUrl)
	proxy := httputil.NewSingleHostReverseProxy(target)
	originalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		originalDirector(req)

		// 重置 body（防止被提前读取为空）
		if req.Body != nil {
			body, err := io.ReadAll(req.Body)
			if err == nil {
				req.Body = io.NopCloser(bytes.NewBuffer(body))
			}
		}

		req.Header.Set("Authorization", "Bearer "+apiKey)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept", "application/json")
		req.Host = target.Host
		allowedHeaders := map[string]bool{
			"Content-Type":  true,
			"Authorization": true,
			"Accept":        true,
			"User-Agent":    true,
		}
		for k := range req.Header {
			if !allowedHeaders[k] {
				delete(req.Header, k)
			}
		}
	}
	return &httpProxyImpl{
		target: target,
		proxy:  proxy,
	}
}

type OpenAIProxy struct {
	*httpProxyImpl
}

func newOpenAIProxy(apiKey, baseUrl string) *OpenAIProxy {
	return &OpenAIProxy{
		httpProxyImpl: newChatProxy(apiKey, baseUrl),
	}
}

type DeepSeekProxy struct {
	*httpProxyImpl
}

func newDeepSeekProxy(apiKey, baseUrl string) *DeepSeekProxy {
	return &DeepSeekProxy{
		httpProxyImpl: newChatProxy(apiKey, baseUrl),
	}
}

type ChatProxyFactory struct{}

func (f *ChatProxyFactory) NewProxies(m map[string]models.Provider) map[string]HttpProxy {
	proxies := make(map[string]HttpProxy)
	for name, provider := range m {
		baseUrl := provider.BaseUrl
		switch name {
		case "openai":
			proxies[name] = newOpenAIProxy(provider.APIKey, baseUrl)
		case "deepseek":
			proxies[name] = newDeepSeekProxy(provider.APIKey, baseUrl)
		}
	}

	return proxies
}
func (f *ChatProxyFactory) NewProxy(providerName string, apiKey, baseUrl string) HttpProxy {
	switch providerName {
	case "openai":
		return newOpenAIProxy(apiKey, baseUrl)
	case "deepseek":
		return newDeepSeekProxy(apiKey, baseUrl)
	default:
		return nil
	}
}
